import paho.mqtt.client as mqtt
import json
class MQTTClient:
    def __init__(self, broker,topic_sub ,topic_pub,port=21883,client_id="", username=None, password=None, keepalive=60, tls=False):
        """
        初始化 MQTT 客户端
        :param broker: MQTT 代理地址
        :param port: MQTT 代理端口 (默认 1883)
        :param client_id: 客户端 ID，默认为空字符串，自动生成
        :param username: 用户名（可选）
        :param password: 密码（可选）
        :param keepalive: 心跳时间，单位秒（默认 60）
        :param tls: 是否使用 TLS/SSL 加密（默认 False）
        """
        self.client = mqtt.Client(client_id)
        self.broker = broker
        self.port = port
        self.keepalive = keepalive
        self.topic_sub = topic_sub
        self.topic_pub = topic_pub

        # 设置用户名和密码
        if username and password:
            self.client.username_pw_set(username, password)

        # 是否启用 TLS/SSL
        if tls:
            self.client.tls_set()

        # 回调函数默认设置为空函数，可自定义
        self.client.on_connect = self._on_connect
        self.client.on_message = self._on_message
        self.client.on_disconnect = self._on_disconnect
        self.client.on_publish = self._on_publish

    def send_json_msg(self, msg):
        """
        parameters:msg:json消息
        returns:无
        functions:在指定的话题上发布消息
        """
        self.client.publish(self.topic_pub, payload="{}".format(msg))

    def _on_connect(self, client, userdata, flags, rc):
        """默认的连接成功回调函数"""
        if rc == 0:
            print("Connected successfully")
        else:
            print(f"Failed to connect, return code {rc}")

    def _on_message(self, client, userdata, msg):
        """默认的接收消息回调函数"""
        print(f"Received message: '{msg.payload.decode()}' on topic '{msg.topic}'")

    def _on_disconnect(self, client, userdata, rc):
        """默认的断开连接回调函数"""
        print("Disconnected")

    def _on_publish(self, client, userdata, mid):
        """默认的消息发布回调函数"""
        # print(f"Message {mid} published")

    def connect(self):
        """连接到 MQTT 代理"""
        self.client.connect(self.broker, self.port, self.keepalive)
        # 订阅主题
        self.client.subscribe(self.topic_sub)
        # 开启接收循环
        self.client.loop_start()

    def disconnect(self):
        """断开与 MQTT 代理的连接"""
        self.client.disconnect()

    def publish(self, topic, message, qos=0, retain=False):
        """
        发布消息到指定主题
        :param topic: 主题名称
        :param message: 消息内容
        :param qos: 服务质量 (0, 1, 2)
        :param retain: 是否保留消息
        """
        self.client.publish(topic, message, qos, retain)

    def subscribe(self, topic, qos=0):
        """
        订阅指定主题
        :param topic: 主题名称
        :param qos: 服务质量 (0, 1, 2)
        """
        self.client.subscribe(topic, qos)

    def set_on_connect(self, callback):
        """设置连接成功回调"""
        self.client.on_connect = callback

    def set_on_message(self, callback):
        """设置消息接收回调"""
        self.client.on_message = callback

    def set_on_disconnect(self, callback):
        """设置断开连接回调"""
        self.client.on_disconnect = callback

    def set_on_publish(self, callback):
        """设置消息发布回调"""
        self.client.on_publish = callback

    def control_device(self, str_key, str_value):
        """
        parameters:str_key:键的字符串
                   str_value:值的字符串
        functions:控制设备
        """
        self.send_json_msg(json.dumps({str_key: str_value}))

    def loop_forever(self):
        """开始阻塞式网络循环"""
        self.client.loop_forever()

    def loop_start(self):
        """开始非阻塞式网络循环"""
        self.client.loop_start()

    def loop_stop(self):
        """停止非阻塞式网络循环"""
        self.client.loop_stop()

    def __del__(self):
        self.client.loop_stop()  # 结束接收循环
        self.client.disconnect()  # 断开连接
